package com.luo.demo.scg.client.ext;

import com.luo.demo.scg.client.config.Oauth2ClientSecurityConfig;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.core.log.LogMessage;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.security.oauth2.client.registration.ClientRegistration;
import org.springframework.security.oauth2.client.registration.ReactiveClientRegistrationRepository;
import org.springframework.security.oauth2.client.web.server.DefaultServerOAuth2AuthorizationRequestResolver;
import org.springframework.security.oauth2.core.endpoint.OAuth2AuthorizationRequest;
import org.springframework.security.web.server.savedrequest.WebSessionServerRequestCache;
import org.springframework.util.CollectionUtils;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.util.Optional;

/**
 * OAuth2 Client Authorization Endpoint /oauth2/authoriztion/{clientRegId} <br/>
 * 请求解析器扩展实现 - 支持提取query参数redirect_uri，用作后续OAuth2认证完成后SCG重定向到该指定redirect_uri。<br/>
 * 适用场景：SPA -> SCG -> SCG返回401 -> SPA重定向到/oauth2/authoriztion/{clientRegId}?redirect_uri=http://spa -> SCG完成OAuth2认证后再重定向回http://spa
 *
 * @author luohq
 * @date 2022-06-13
 * @see DefaultServerOAuth2AuthorizationRequestResolver
 * @see WebSessionServerRequestCache
 * @see Oauth2ClientSecurityConfig
 */
public class SaveRequestServerOAuth2AuthorizationRequestResolver extends DefaultServerOAuth2AuthorizationRequestResolver {

    private static final Log logger = LogFactory.getLog(SaveRequestServerOAuth2AuthorizationRequestResolver.class);

    /**
     * redirect uri参数名称
     */
    private static final String PARAM_REDIRECT_URI = "redirect_uri";

    /**
     * WebSession对应的saveRequest属性名
     * 完全沿用（兼容）WebSessionServerRequestCache定义
     */
    private static final String DEFAULT_SAVED_REQUEST_ATTR = "SPRING_SECURITY_SAVED_REQUEST";
    private String sessionAttrName = DEFAULT_SAVED_REQUEST_ATTR;


    /**
     * Creates a new instance
     *
     * @param clientRegistrationRepository the repository to resolve the
     *                                     {@link ClientRegistration}
     */
    public SaveRequestServerOAuth2AuthorizationRequestResolver(
            ReactiveClientRegistrationRepository clientRegistrationRepository) {
        super(clientRegistrationRepository);
    }

    @Override
    public Mono<OAuth2AuthorizationRequest> resolve(ServerWebExchange exchange) {
        return super.resolve(exchange)
                .doOnNext(oAuth2AuthorizationRequest -> {
                    //获取query参数redirect_uri
                    Optional.ofNullable(exchange.getRequest())
                            .map(ServerHttpRequest::getQueryParams)
                            .map(queryParams -> queryParams.get(PARAM_REDIRECT_URI))
                            .filter(redirectUris -> !CollectionUtils.isEmpty(redirectUris))
                            .map(redirectUris -> redirectUris.get(0))
                            .ifPresent(redirectUri -> {
                                //若redirect_uri非空，则覆盖Session中的SPRING_SECURITY_SAVED_REQUEST为redirect_uri
                                //即后续认证成功后可重定向回SPA指定页面
                                exchange.getSession().subscribe(webSession -> {
                                    webSession.getAttributes().put(this.sessionAttrName, redirectUri);
                                    logger.debug(LogMessage.format("SCG OAuth2 authorization endpoint queryParam redirect_uri added to WebSession: '%s'", redirectUri));
                                });
                            });
                });
    }

}

